import warnings
warnings.filterwarnings("ignore")

import os
import tensorflow as tf

LOGDIR = 'log'

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
tf.compat.v1.logging.set_verbosity(40)

from tensorflow.examples.tutorials.mnist import input_data

tf.set_random_seed(777)

LOGDIR = "log"

def conv_layer(input, in_size, out_size, padding, name="conv1"):
    with tf.name_scope(name):
        w = tf.Variable(tf.truncated_normal([5, 5, in_size, out_size], stddev=0.1), name="w")
        b = tf.Variable(tf.constant(0.1, shape=[out_size]), name="b")
        conv = tf.nn.conv2d(input, w, strides=[1, 1, 1, 1], padding=padding)
        a = tf.nn.relu(conv + b)
        return a

def fc_layer(input, in_size, out_size, name):
    with tf.name_scope(name):
        w = tf.Variable(tf.truncated_normal([in_size, out_size], stddev=0.1), name="w")
        b = tf.Variable(tf.constant(0.1, shape=[out_size]), name="b")
        z = tf.matmul(input, w) + b
        return z

def cnn_model(x):
    conv1 = conv_layer(x, 1, 6, padding="SAME", name="conv1")
    conv1_pool = tf.nn.max_pool2d(conv1, [1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME")
    conv2 = conv_layer(conv1_pool, 6, 16, padding="VALID", name="conv2")
    conv2_pool = tf.nn.max_pool2d(conv2, [1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME")
    flatten = tf.reshape(conv2_pool, [-1, 5 * 5 * 16])

    fc3 = tf.nn.relu(fc_layer(flatten, 5 * 5 * 16, 120, "fc3"))
    fc4 = tf.nn.relu(fc_layer(fc3, 120, 84, "fc4"))
    logits = fc_layer(fc4, 84, 10, "logits")
    return logits

def mnist():
    with tf.name_scope("input"):
        x = tf.placeholder(tf.float32, shape=[None, 784], name="x")
        x_image = tf.reshape(x, [-1, 28, 28, 1])
        y = tf.placeholder(tf.float32, shape=[None, 10], name="y")

    with tf.name_scope("cnn_model"):
        logits = cnn_model(x_image)

    with tf.name_scope("loss"):
        loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
            logits=logits, labels=y))

    # optimize
    with tf.name_scope("train"):
        train_op = tf.train.AdamOptimizer(learning_rate=1e-4).minimize(loss)

    with tf.name_scope("accuracy"):
        correct = tf.equal(tf.argmax(logits, 1), tf.argmax(y, 1))
        accuracy = tf.reduce_mean(tf.cast(correct, tf.float32))

    # data sets
    mnist = input_data.read_data_sets("../MNIST_data", one_hot=True)

    # launch a session
    sess = tf.Session()
    sess.run(tf.global_variables_initializer())

    writer = tf.summary.FileWriter(LOGDIR)
    writer.add_graph(sess.graph)

    training_epochs = 2
    batch_size = 100
    for epoch in range(training_epochs):
        avg_cost = 0
        total_batch = int(mnist.train.num_examples / batch_size)
        for i in range(total_batch):
            xs, ys = mnist.train.next_batch(batch_size)
            cost_val, _ = sess.run([loss, train_op], feed_dict={x: xs, y: ys})
            avg_cost += cost_val
        avg_cost = avg_cost /total_batch
        print("epoch:", epoch+1, "cost: ", avg_cost)

    # accuracy
    test_acc = sess.run(accuracy, feed_dict={x: mnist.test.images, y: mnist.test.labels})
    print("test acc: ", test_acc)

if __name__ == '__main__':
    mnist()
